"""基于论文[1] Christian J. A. , Derksen H. , Watkins R. .Lunar crater identification in digital Images[J/OL].J. Astronaut. Sci.,2021,68(4):1056-1144

"""

from .base import TriadHierarchyBuilder
from utils.ellipse import radius_ellipse
import numpy as np


class ChristianHierarchyBuilder(TriadHierarchyBuilder):
    def invariant(self, i, j, k, Q1, Q2, Q3, th=0) -> tuple[np.ndarray, np.ndarray]:
        """
        TODO: 排除直径相接近的情况，可以不用排除共线、包含与相交的情况。
        """
        N = Q3.shape[0]
        Q1 = Q1[None].repeat(N, axis=0)
        Q2 = Q2[None].repeat(N, axis=0)
        d1 = np.mean(radius_ellipse(Q1), axis=0)
        d2 = np.mean(radius_ellipse(Q2), axis=0)
        d3 = np.mean(radius_ellipse(Q3), axis=0)
        # 按直径从小到大排序
        D = np.array((d1, d2, d3))
        # 任取两个直径计算比值，当比值接近1的数超过2个时，排除
        remain_ind = np.abs(D[[1, 2]] / D[0, None] - 1) > th
        remain_ind &= np.abs(D[[0, 2]] / D[1, None] - 1) > th
        remain_ind &= np.abs(D[[0, 1]] / D[2, None] - 1) > th
        remain_ind = remain_ind.all(axis=0)
        ind = np.argsort((d1, d2, d3), axis=0)
        Q = np.array((Q1, Q2, Q3))[ind, np.arange(N)]
        # 计算不变量
        I1 = np.trace(Q[0] @ np.linalg.inv(Q[1]), axis1=1, axis2=2) * np.trace(
            Q[1] @ np.linalg.inv(Q[0]), axis1=1, axis2=2
        )
        I2 = np.trace(Q[1] @ np.linalg.inv(Q[2]), axis1=1, axis2=2) * np.trace(
            Q[2] @ np.linalg.inv(Q[1]), axis1=1, axis2=2
        )
        I3 = np.trace(Q[2] @ np.linalg.inv(Q[0]), axis1=1, axis2=2) * np.trace(
            Q[0] @ np.linalg.inv(Q[2]), axis1=1, axis2=2
        )
        # 每行是一个3x6的点列，每三个元素代表一个点
        I = np.array((I1, I2, I3))
        ijk = np.vstack((np.ones_like(k) * i, np.ones_like(k) * j, k))
        # 按直径从小到大排序
        ijk = np.take_along_axis(ijk, ind, axis=0)[:, remain_ind]
        return ijk, I[:, remain_ind]
